package com.lnn.springBase.nc;

import ucar.ma2.Array;
import ucar.ma2.DataType;

public class ArrayFactory {
    // float 类型支持
    public static Array makeArray(DataType targetType, int[] shape, float[] data) {
        validateShape(shape, data.length, 1);
        return Array.factory(targetType, shape, data);
    }

    public static Array makeArray(DataType targetType, int[] shape, float[][] data) {
        validateShape(shape, data.length, data[0].length);
        return Array.factory(targetType, shape, flatten2D(data));
    }

    public static Array makeArray(DataType targetType, int[] shape, float[][][] data) {
        validate3DShape(shape, data);
        return Array.factory(targetType, shape, flatten3D(data));
    }

    // double 类型支持
    public static Array makeArray(DataType targetType, int[] shape, double[] data) {
        validateShape(shape, data.length, 1);
        return Array.factory(targetType, shape, data);
    }

    public static Array makeArray(DataType targetType, int[] shape, double[][] data) {
        validateShape(shape, data.length, data[0].length);
        return Array.factory(targetType, shape, flatten2D(data));
    }

    public static Array makeArray(DataType targetType, int[] shape, double[][][] data) {
        validate3DShape(shape, data);
        return Array.factory(targetType, shape, flatten3D(data));
    }

    // int 类型支持
    public static Array makeArray(DataType targetType, int[] shape, int[] data) {
        validateShape(shape, data.length, 1);
        return Array.factory(targetType, shape, data);
    }

    public static Array makeArray(DataType targetType, int[] shape, int[][] data) {
        validateShape(shape, data.length, data[0].length);
        return Array.factory(targetType, shape, flatten2D(data));
    }

    public static Array makeArray(DataType targetType, int[] shape, int[][][] data) {
        validate3DShape(shape, data);
        return Array.factory(targetType, shape, flatten3D(data));
    }

    // 维度校验方法
    private static void validateShape(int[] shape, int... actualDims) {
        if (shape.length != actualDims.length) {
            throw new IllegalArgumentException("维度不匹配，预期维度数："
                    + shape.length + "，实际维度数：" + actualDims.length);
        }

        for (int i = 0; i < shape.length; i++) {
            if (shape[i] != actualDims[i]) {
                throw new IllegalArgumentException(String.format(
                        "第 %d 维度长度不匹配，预期：%d，实际：%d",
                        i + 1, shape[i], actualDims[i]
                ));
            }
        }
    }

    // 三维数组专用校验
    private static void validate3DShape(int[] shape, Object data) {
        if (shape.length != 3) throw new IllegalArgumentException("需要三维形状参数");

        float[][][] arr3D = (float[][][]) data;
        int z = arr3D.length;
        int y = arr3D[0].length;
        int x = arr3D[0][0].length;

        if (shape[0] != z || shape[1] != y || shape[2] != x) {
            throw new IllegalArgumentException(String.format(
                    "三维数组尺寸不匹配，预期：%dx%dx%d，实际：%dx%dx%d",
                    shape[0], shape[1], shape[2], z, y, x
            ));
        }
    }

    // 多维数组扁平化工具
    private static float[] flatten3D(float[][][] data) {
        int total = data.length * data[0].length * data[0][0].length;
        float[] flat = new float[total];
        int index = 0;
        for (float[][] layer : data) {
            for (float[] row : layer) {
                System.arraycopy(row, 0, flat, index, row.length);
                index += row.length;
            }
        }
        return flat;
    }

    private static double[] flatten3D(double[][][] data) {
        // 实现与 float 版本类似...
        int total = data.length * data[0].length * data[0][0].length;
        double[] flat = new double[total];
        int index = 0;
        for (double[][] layer : data) {
            for (double[] row : layer) {
                System.arraycopy(row, 0, flat, index, row.length);
                index += row.length;
            }
        }
        return flat;
    }

    private static int[] flatten3D(int[][][] data) {
        int total = data.length * data[0].length * data[0][0].length;
        int[] flat = new int[total];
        int index = 0;
        for (int[][] layer : data) {
            for (int[] row : layer) {
                System.arraycopy(row, 0, flat, index, row.length);
                index += row.length;
            }
        }
        return flat;
    }

    private static int[] flatten2D(int[][] data) {
        int total = data.length * data[0].length;
        int[] flat = new int[total];
        int index = 0;
        for (int[] row : data) {
            System.arraycopy(row, 0, flat, index, row.length);
            index += row.length;
        }
        return flat;
    }

    private static double[] flatten2D(double[][] data) {
        int total = data.length * data[0].length;
        double[] flat = new double[total];
        int index = 0;
        for (double[] row : data) {
            System.arraycopy(row, 0, flat, index, row.length);
            index += row.length;
        }
        return flat;
    }

    private static float[] flatten2D(float[][] data) {
        int total = data.length * data[0].length;
        float[] flat = new float[total];
        int index = 0;
        for (float[] row : data) {
            System.arraycopy(row, 0, flat, index, row.length);
            index += row.length;
        }
        return flat;
    }
}
